import pandas as pd
import requests
from pathlib import Path
import threading
from queue import Queue, Empty
import time
from concurrent.futures import ThreadPoolExecutor
import logging
from tqdm import tqdm
import os
import traceback

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('download_error.log'),
        logging.StreamHandler()
    ]
)

# 配置
NUM_THREADS = 20
TIMEOUT = 10  # 请求超时时间（秒）
RETRY_TIMES = 3  # 下载重试次数

class ImageDownloader:
    def __init__(self, csv_path, save_root, num_threads=NUM_THREADS):
        self.csv_path = Path(csv_path)
        self.save_root = Path(save_root)
        self.num_threads = num_threads
        self.queue = Queue()
        self.failed_downloads = Queue()
        self.session = requests.Session()
        self.pbar = None
        
        # 创建保存目录
        self.save_root.mkdir(parents=True, exist_ok=True)
        
    def load_data(self):
        """加载CSV数据并填充队列"""
        df = pd.read_csv(self.csv_path)
        total_images = len(df)
        self.pbar = tqdm(total=total_images, desc="Downloading images")
        
        for _, row in df.iterrows():
            try:
                goods_id = str(row['goods_id'])
                image_id = str(row['image_id'])
                image_url = row['image_url']
                
                if pd.isna(image_url) or pd.isna(goods_id) or pd.isna(image_id):
                    logging.warning(f"Invalid data: goods_id={goods_id}, image_id={image_id}, url={image_url}")
                    continue
                    
                self.queue.put({
                    'goods_id': goods_id,
                    'image_id': image_id,
                    'url': image_url
                })
            except Exception as e:
                logging.error(f"Error processing row: {row}\nError: {str(e)}")
                continue

    def download_image(self, url, save_path):
        """下载单个图片，带重试机制"""
        for attempt in range(RETRY_TIMES):
            try:
                response = self.session.get(url, timeout=TIMEOUT)
                response.raise_for_status()
                
                # 检查是否是图片内容
                content_type = response.headers.get('content-type', '')
                if not content_type.startswith('image/'):
                    raise ValueError(f"Not an image: {content_type}")
                
                # 保存图片
                save_path.parent.mkdir(parents=True, exist_ok=True)
                with open(save_path, 'wb') as f:
                    f.write(response.content)
                return True
                
            except requests.RequestException as e:
                if attempt == RETRY_TIMES - 1:  # 最后一次重试也失败
                    logging.error(f"Failed to download {url} after {RETRY_TIMES} attempts: {str(e)}")
                    return False
                time.sleep(1)  # 重试前等待
            except Exception as e:
                logging.error(f"Unexpected error downloading {url}: {str(e)}")
                return False
                
    def download_worker(self):
        """工作线程函数"""
        while True:
            try:
                # 获取下一个任务
                item = self.queue.get(block=True, timeout=5)  # 5秒超时
                
                # 构建保存路径
                filename = f"{item['goods_id']}_{item['image_id']}.jpg"
                save_path = self.save_root / filename
                
                # 如果文件已存在，跳过下载
                if save_path.exists():
                    self.pbar.update(1)
                    self.queue.task_done()
                    continue
                
                # 下载图片
                success = self.download_image(item['url'], save_path)
                
                if not success:
                    self.failed_downloads.put(item)
                
                self.pbar.update(1)
                self.queue.task_done()
                
            except Empty:
                break  # 队列为空，退出线程
            except Exception as e:
                logging.error(f"Worker error: {str(e)}\n{traceback.format_exc()}")
                continue
                
    def start(self):
        """开始下载过程"""
        try:
            logging.info("Starting download process...")
            
            # 加载数据
            self.load_data()
            total_images = self.queue.qsize()
            logging.info(f"Loaded {total_images} images to download")
            
            # 创建线程池
            threads = []
            for _ in range(self.num_threads):
                thread = threading.Thread(target=self.download_worker)
                thread.start()
                threads.append(thread)
            
            # 等待所有线程完成
            for thread in threads:
                thread.join()
            
            # 处理失败的下载
            failed_count = self.failed_downloads.qsize()
            if failed_count > 0:
                logging.warning(f"Failed to download {failed_count} images")
                # 将失败的下载记录到文件
                with open('temp/failed_downloads.txt', 'w') as f:
                    while not self.failed_downloads.empty():
                        item = self.failed_downloads.get()
                        f.write(f"{item['goods_id']},{item['image_id']},{item['url']}\n")
            
            logging.info("Download process completed!")
            
        except Exception as e:
            logging.error(f"Error in download process: {str(e)}\n{traceback.format_exc()}")
        finally:
            if self.pbar:
                self.pbar.close()

def main():
    # 配置参数
    csv_path = 'data/goods_images_filtered_10w.csv'  # 使用之前分割好的Q表
    save_root = '/home/tfj/datasets/image_retri10k/eval_images_o2'  # 图片保存目录
    os.makedirs(save_root, exist_ok=True)
    # 创建下载器实例并开始下载
    downloader = ImageDownloader(csv_path, save_root, NUM_THREADS)
    downloader.start()

if __name__ == "__main__":
    main() 